import cv2
import numpy as np
import matplotlib.pyplot as plt
from skimage import feature, exposure, transform, morphology
from scipy.signal import find_peaks
from skimage.morphology import footprint_rectangle

def show_img(img):
    plt.imshow(img)
    plt.axis('off')
    plt.show()

def image_preprocessing(image_path):
    """图片预处理"""
    # 1. 读取图像并利用高斯模糊去噪
    img = cv2.imread(image_path, 0)
    if img is None:
        print(f"错误:无法加载图像，请检查路径:{image_path}")
        return
    blurred = cv2.GaussianBlur(img, (3, 3), 0)

    # 2. 自适应对比度增强
    p2, p98 = np.percentile(blurred, (4, 98)) # 4% 的像素值被视为异常暗，2% 的像素值被视为异常亮,排除
    enhanced = exposure.rescale_intensity(blurred, in_range=(p2, p98)) # 对非极端部分提高对比度

    # 3. 自适应边缘检测:检测图像明暗变化梯度，确定边缘（线条）位置
    sobel_x = cv2.Sobel(enhanced, cv2.CV_64F, 1, 0, ksize=3) # 水平梯度检测
    sobel_y = cv2.Sobel(enhanced, cv2.CV_64F, 0, 1, ksize=3) # 垂直梯度检测
    # 两方向梯度平方相加并开根号:计算梯度幅值，表示图像中每个像素点处亮度变化的强度，幅值大的位置对应图像中的边缘
    gradient_mag = np.sqrt(sobel_x ** 2 + sobel_y ** 2)

    grad_mean = np.mean(gradient_mag) # 计算整个图像梯度幅值的平均值
    # 更具图像梯度幅值确定sigma值
    if grad_mean > 233:
        sigma = 5
    else:
        sigma = 3
    # 提高亮度
    edges = feature.canny(enhanced, sigma=sigma)
    # edges = morphology.closing(edges, footprint_rectangle((3, 3))) # 闭合孔隙，减少断线
    return [img, enhanced, edges]


def measure_stripe_spacing(image_path, scale_bar_pixels, scale_bar_actual, manual_mode=False):
    """测量电镜图中条纹间隔的主函数，支持自动和手动模式"""

    img, enhanced, edges = image_preprocessing(image_path=image_path)
    # 4. 显示中间结果:三张图进行对比
    plt.figure(figsize=(12, 8))
    plt.subplot(221), plt.imshow(img, cmap='gray'), plt.title('原图')
    plt.subplot(222), plt.imshow(enhanced, cmap='gray'), plt.title('对比度增强')
    plt.subplot(223), plt.imshow(edges, cmap='gray'), plt.title('边缘检测结果')
    plt.tight_layout()
    plt.show()

    # 5. 手动框选条纹区域
    print("请在弹出窗口中框选包含条纹的矩形区域（尽量包含多条完整条纹），按回车确认:")
    roi = cv2.selectROI("框选条纹区域", img) # 打开交互窗口，显示原图
    x, y, w, h = roi # 获得左上角坐标和图像高度和宽度
    roi_img = img[y:y + h, x:x + w] # 提取原图像区域
    roi_edges = edges[y:y + h, x:x + w] # 提取edges图像区域
    cv2.destroyAllWindows() # 关闭交互窗口

    # 6. 检测线条方向，实现图像旋转
    if not manual_mode: # 如果默认非手动
        print("正在自动检测条纹方向...")
        roi_edges_uint8 = roi_edges.astype(np.uint8) * 255 # 图片类型转换，后续需要

        # 霍夫变换检测直线，返回张量shape为(ddd, 1, 4)，每(1, 1, 4)储存两个端点:[[[x1, y1, x2, y2]]]
        lines = cv2.HoughLinesP(
            roi_edges_uint8,
            rho=1, # 距离分辨率（像素），表示霍夫空间中 ρ 的步长。
            theta=np.pi / 180, # 角度分辨率（弧度），表示霍夫空间中 θ 的步长，这里设为 1 度。
            threshold=10,  # 累加器阈值，只有投票数超过该值的线段才会被保留。较低的阈值会增加检测到的直线数量。
            minLineLength=15, # 最小线段长度（像素），短于该长度的线段会被过滤。
            maxLineGap=3 # 线段最大间隔（像素），若两条线段间距小于此值，会被视为同一条直线。
        )
        print(lines.shape)

        # 绘制检测到的线段
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0] # 获得线段两端点的坐标
                # 绘制线段
                plt.plot((x1, x2), (-y1, -y2), 'k-', linewidth=2)  # 黑色线段
                # 绘制端点，使用不同颜色
                plt.plot(x1, -y1, 'ro', markersize=3)  # 红色起点
                plt.plot(x2, -y2, 'bo', markersize=3)  # 蓝色终点
                plt.title('分别绘制线段和端点')
                plt.grid(True)
            plt.show()

        # 可视化:对比三张图
        img_copied = roi_img.copy() # 复制一份剪切图原图，用做画布

        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                cv2.line(img_copied, (x1, y1), (x2, y2), (0, 255, 0), 2)
        plt.figure(figsize=(15, 5))
        plt.subplot(131), plt.imshow(cv2.cvtColor(roi_img, cv2.COLOR_BGR2RGB))
        plt.title('原始图像'), plt.axis('off')
        plt.subplot(132), plt.imshow(edges, cmap='gray')
        plt.title('边缘图'), plt.axis('off')
        plt.subplot(133), plt.imshow(cv2.cvtColor(img_copied, cv2.COLOR_BGR2RGB))
        plt.title('检测结果'), plt.axis('off')
        plt.tight_layout()
        plt.show()

        # 在原图上绘制检测到的直线
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                cv2.line(img_copied, (x1, y1), (x2, y2), (0, 255, 0), 1)
            plt.figure(), plt.imshow(img_copied, cmap='gray'), plt.title('霍夫检测结果')
            plt.show()

        # 傅里叶变换:傅里叶变换是一种将时域信号转换为频域信号的数学工具
        f = np.fft.fft2(roi_img) # 二维傅里叶变换
        fshift = np.fft.fftshift(f) # 将低频分量移动到频谱数组的中心
        magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1) # 转换值，便于后续处理，20可更改

        # 可视化频谱
        plt.figure(figsize=(15, 5))
        plt.subplot(131), plt.imshow(magnitude_spectrum, cmap='gray')
        plt.title('对数幅值谱'), plt.axis('off')
        plt.subplot(132), plt.imshow(magnitude_spectrum, cmap='jet')
        plt.title('对数幅值谱'), plt.colorbar(label='频率')
        plt.tight_layout()
        plt.show()

        # 创建环形掩码，用于过滤傅里叶频谱中的特定频率区域，保留中频信息
        # 获取频谱形状找到中心
        h, w = magnitude_spectrum.shape # 获取频谱长宽
        center = (w // 2, h // 2) # 除以2取整，找到中心
        print(center)
        # 构建一个结构相同的矩阵，存储每个像素的列坐标和行坐标。[0]为行，[1]为高
        y, x = np.indices((h, w))

        r = np.sqrt((x - center[0]) ** 2 + (y - center[1]) ** 2) # 计算每一点到中心的距离
        inner_radius = 15 # 内部排除半径，环内半径
        outer_radius = min(h, w) // 4 # 外部排除半径，环外半径
        r_mask = (r > inner_radius) & (r < outer_radius) # 布尔型数组，环内部分值为True
        # 可视化掩码
        plt.figure(figsize=(15, 5))
        plt.subplot(131), plt.imshow(r, cmap='jet')
        plt.title('到中心的距离'), plt.colorbar(label='半径')
        plt.scatter(center[0], center[1], color='red', s=50)  # 标记中心
        plt.subplot(132), plt.imshow(r_mask, cmap='gray')
        plt.title('环形掩码'), plt.axis('off')
        plt.subplot(133), plt.imshow(np.ones_like(r_mask), cmap='gray')
        plt.title('掩码叠加示意图'), plt.axis('off')
        circle1 = plt.Circle(center, inner_radius, color='r', fill=False)
        circle2 = plt.Circle(center, outer_radius, color='r', fill=False)
        plt.gca().add_patch(circle1)
        plt.gca().add_patch(circle2)
        plt.tight_layout()
        plt.show()

        print(x, y)
        # 计算每一点相对中心的角度(以中心为原点相对x轴)，并换算为角度制，构建一个矩阵
        angles = np.arctan2(y - center[1], x - center[0]) * 180 / np.pi
        print('angles:', angles)
        # 角度范围扩展到±180°
        print('angles_masked:', angles[r_mask])

        weights = magnitude_spectrum[r_mask].flatten() # 频率权重
        # 统计数组中值的分布情况，仅考虑环形掩码内的点，将角度范围(-180°, 180°)分成720个区间
        # hist:每个区间内的元素数量（整数数组）;每个区间边界取小不取大，如[1, 2)
        # bins:区间边界数组（长度比 hist 多 1）
        angles_0_180 = angles % 180  # 转换到0-180度范围
        hist, bins = np.histogram(angles_0_180[r_mask], bins=180, range=(0, 180), weights=weights)
        print('hist:', hist)
        print('bins:', bins)

        # 可视化
        plt.figure(figsize=(10, 5))
        plt.bar(bins[:-1], hist, width=0.5, align='edge', alpha=0.7)
        plt.axvline(x=3, color='r', linestyle='--', label='真实方向')
        plt.title('角度分布直方图')
        plt.xlabel('角度 (度)')
        plt.ylabel('频数')
        plt.legend()
        plt.grid(True)
        plt.show()

        # argmax找到分布最多的方向的索引，从bins中提取出角度
        angle_fourier = bins[np.argmax(hist)]
        print('max_hist:', np.argmax(hist))
        print('angle_fourier:', angle_fourier)

        # 旋转角度选择
        angle_hough = 0

        if lines is not None and len(lines) > 3:
            angles_hough = []
            filtered_lines = []
            for line in lines:
                x1, y1, x2, y2 = line[0]
                dx = x2 - x1
                dy = -(y2 - y1) # 坐标轴转换
                angle = np.degrees(np.arctan2(dy, dx))
                if angle > 90:
                    angle -= 180
                elif angle < -90:
                    angle += 180
                angles_hough.append(angle)
                filtered_lines.append(line)

            # 可视化
            plt.figure(figsize=(15, 5))
            # 原始直线
            plt.subplot(121)
            plt.title('直线检测结果')
            plt.xlim(0, 500), plt.ylim(400, 0)  # 注意:图像坐标系y轴向下
            for line in lines:
                x1, y1, x2, y2 = line[0]
                plt.plot([x1, x2], [y1, y2], 'b-', linewidth=2)
            plt.figure(figsize=(10, 5))
            plt.imshow(magnitude_spectrum, cmap='gray')
            plt.quiver(center[0], center[1],
                       np.cos(np.deg2rad(angle_fourier)),
                       np.sin(np.deg2rad(angle_fourier)),
                       color='red', scale=50)
            plt.title(f'检测到的法线方向: {angle_fourier:.2f}°')
            plt.show()

            if angles_hough:
                def circular_median(angles):
                    angles_rad = np.deg2rad(angles)
                    median_sin = np.median(np.sin(angles_rad))
                    median_cos = np.median(np.cos(angles_rad))
                    return np.rad2deg(np.arctan2(median_sin, median_cos))

                angle_hough = circular_median(angles_hough)
                rotation_angle = -angle_hough
                method = "霍夫变换"
            else:
                stripe_direction = (angle_fourier + 90) % 180

                if stripe_direction > 90:
                    stripe_direction -= 180
                rotation_angle = -stripe_direction
                method = "傅里叶变换（无有效大角度）"
        else:
            stripe_direction = (angle_fourier + 90) % 180

            if stripe_direction > 90:
                stripe_direction -= 180
            rotation_angle = -stripe_direction
            method = "傅里叶变换（无有效直线）"

        print(f"自动检测结果:条纹方向 {rotation_angle:.2f} 度 ({method})")
        print(f"霍夫角度参考:{angle_hough:.2f}度（若有），傅里叶角度:{angle_fourier:.2f}度")

        # 手动验证提示
        manual_check = input("条纹是否水平？需要手动调整角度吗？(y/n):")
        if manual_check.lower() == 'y':
            rotation_angle = float(input("请输入正确角度（度）:"))
    else:
        rotation_angle = float(input("请输入条纹的倾斜角度（度）:"))

    # 7. 旋转图像
    rotated_roi = transform.rotate(roi_img, -rotation_angle, mode='constant')
    rotated_edges = transform.rotate(roi_edges_uint8, -rotation_angle, mode='constant')
    # 可视化


    # 8. 投影方向选择
    print("请选择投影方向:1) 垂直 2) 水平 (默认垂直)")
    choice = input("输入数字(1/2):")
    if choice == '2':
        projection_axis = 0
        projection_dir = "水平"
    else:
        projection_axis = 1
        projection_dir = "垂直"

    print(f"选择{projection_dir}投影分析条纹间隔")
    projection = np.sum(rotated_edges, axis=projection_axis)
    # 可视化
    plt.figure(figsize=(10, 5))
    plt.subplot(121), plt.imshow(rotated_edges, cmap='gray')
    plt.title('校正后的垂直条纹'), plt.axis('off')
    plt.subplot(122), plt.plot(projection)
    plt.title('垂直投影曲线'), plt.xlabel('位置'), plt.ylabel('像素和')
    plt.grid(True)
    plt.tight_layout()
    plt.show()

    # 9. 优化峰值检测（减少误判）
    proj_mean = np.mean(projection)
    proj_std = np.std(projection)
    proj_max = np.max(projection)
    proj_min = np.min(projection)
    proj_range = proj_max - proj_min  # 投影对比度

    # 动态阈值（根据对比度调整）
    if proj_range > 50:
        threshold = proj_mean + 0.8 * proj_std
    elif proj_range > 30:
        threshold = proj_mean + 0.6 * proj_std
    else:
        threshold = proj_mean + 0.4 * proj_std

    # 增加峰值突出度过滤
    peaks, properties = find_peaks(
        projection,
        height=threshold,
        distance=10,
        prominence=10
    )

    if len(peaks) >= 2:
        pixel_spacings = np.diff(peaks)
        avg_pixel_spacing = np.mean(pixel_spacings)
        actual_spacing = avg_pixel_spacing * (scale_bar_actual / scale_bar_pixels)

        print(f"检测到 {len(peaks)} 个条纹，平均像素间隔:{avg_pixel_spacing:.2f} 像素")
        print(f"实际物理间隔:{actual_spacing:.2f} nm")

        plt.figure(figsize=(10, 6))
        plt.subplot(211), plt.imshow(rotated_roi, cmap='gray')
        plt.title(f'旋转后ROI - 间隔:{actual_spacing:.2f} nm')
        plt.subplot(212), plt.plot(projection)
        plt.axhline(y=threshold, color='r', linestyle='--', label='阈值')
        plt.plot(peaks, projection[peaks], 'ro', label='条纹峰值')
        plt.title(f'{projection_dir}投影结果')
        plt.legend()
        plt.tight_layout()
        plt.show()

        if len(peaks) < 5:
            print("警告:检测到的条纹数量较少，建议手动调整ROI或角度")
    else:
        print("错误:未检测到足够条纹，请调整ROI或参数")
        print("正在尝试降低阈值...")
        threshold = proj_mean + 0.2 * proj_std
        peaks, _ = find_peaks(projection, height=threshold, distance=10)
        if len(peaks) >= 2:
            pixel_spacings = np.diff(peaks)
            avg_pixel_spacing = np.mean(pixel_spacings)
            actual_spacing = avg_pixel_spacing * (scale_bar_actual / scale_bar_pixels)
            print(f"降低阈值后检测到 {len(peaks)} 个条纹，实际间隔:{actual_spacing:.2f} nm")
        else:
            print("仍无法检测，建议使用手动模式或重新框选ROI")


# 运行示例（请修改为实际参数）
if __name__ == "__main__":
    image_path = "../img/new/0006.tif"  # 图像路径
    scale_bar_pixels = 50  # 比例尺像素长度（手动测量）
    scale_bar_actual = 10 # 比例尺实际长度(nm)
    manual_mode = False  # 是否手动输入角度
    plt.rc('font', family='YouYuan')
    measure_stripe_spacing(image_path, scale_bar_pixels, scale_bar_actual, manual_mode)